import os
import sys
import time
import math
from matplotlib import pyplot as plt
from config_manager import VisualizationConfig


class VisualizationManager:
    def __init__(self, config: VisualizationConfig, model_name="model"):
        self.config = config
        self.metrics = config.metrics
        self.show_data = {metric: [] for metric in self.metrics}
        self.general_key = ["train_time", "val_time"]
        self.general_data = {}
        self.visual_dir = os.path.join(self.config.save_dir, model_name)
        os.makedirs(self.visual_dir, exist_ok=True)
        plt.style.use(config.style)

        # 检测是否在 IPython 环境（Jupyter/Colab/VSCode Notebook）
        self._have_ipython = False
        try:
            from IPython.display import clear_output  # type: ignore
            self._ipython_clear = clear_output
            self._have_ipython = True
        except Exception:
            self._have_ipython = False

        # 在 Windows 上尝试启用 ANSI 转义支持
        if os.name == "nt":
            self._enable_windows_ansi()

    # -----------------------------
    # 尝试在 Windows 上启用 ANSI 转义 (VT processing)
    # -----------------------------
    def _enable_windows_ansi(self):
        if os.name != "nt":
            return
        try:
            import ctypes
            kernel32 = ctypes.windll.kernel32
            handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
            mode = ctypes.c_uint()
            if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
                new_mode = mode.value | 0x0004  # ENABLE_VIRTUAL_TERMINAL_PROCESSING
                kernel32.SetConsoleMode(handle, new_mode)
        except Exception:
            pass

    # ---------------------------------
    # 更新指标数据
    # ---------------------------------
    def update(self, total_epoches: int = None, epoch: int = None, **kwargs):
        if not self.config.enable:
            return

        for metric in self.metrics:
            if metric in kwargs["metrics"]:
                self.show_data[metric].append(kwargs["metrics"][metric])

        for key in self.general_key:
            if key in kwargs:
                self.general_data[key] = kwargs[key]

        if epoch is None:
            epoch = len(self.show_data[self.metrics[0]])

        if total_epoches is not None:
            self._draw_progress_bar(total_epoches, epoch)

        if self.config.show_terminal_plot:
            self._draw_terminal()

        if self.config.save_plots:
            self._save_plot(epoch=epoch)

    # ---------------------------------
    # 改进后的终端 ASCII 图显示（固定宽度 + 滚动显示 + 平滑纵向）
    def _draw_terminal(self):
        # self._clear_terminal()

        term_h = max(3, int(self.config.terminal_height))
        term_w = max(10, int(self.config.terminal_width))
        char = self.config.terminal_char

        if not hasattr(self, "_minmax_state"):
            self._minmax_state = {metric: [None, None] for metric in self.metrics}

        for metric in self.metrics:
            values = self.show_data[metric]
            if not values:
                continue

            current_min, current_max = min(values), max(values)

            display_values = self.get_show_values(values, term_w)
            n_points = len(display_values)
            h_len = min(term_h, n_points)
            if h_len <= 0:
                continue

            x_indices = list(range(n_points))
            if any(math.isnan(v) for v in display_values):
                # 打印警告（可删除）
                print(f"[警告] {metric} 中存在 NaN 数据，已忽略并设为 0")
                # 将 NaN 替换为最小值（或 0）
                display_values = [0 if math.isnan(v) else v for v in display_values]

            if current_max == current_min:
                scaled = [1 for _ in display_values]
            else:
                scaled = []
                for v in display_values:
                    # double-check NaN
                    if math.isnan(v):
                        scaled.append(0)
                        continue

                    # 归一化到 0~1
                    norm = (v - current_min) / (current_max - current_min)
                    # 保证 norm 在合法范围内
                    norm = max(0.0, min(1.0, norm))
                    scaled.append(int(norm * (h_len - 1)))

            canvas = [[" " for _ in range(n_points)] for _ in range(h_len)]
            for xi, yi in zip(x_indices, scaled):
                y = h_len - 1 - yi
                y = max(0, min(h_len - 1, y))
                canvas[y][xi] = char

            print(f"{metric.upper()}:")
            for row in canvas:
                print("".join(row))
            print(f"Latest {metric}: {values[-1]:.4f}  (min={current_min:.4f}, max={current_max:.4f})")

        print_str = ""
        for key, value in self.general_data.items():
            print_str += f"{key}: {value}    "
        if print_str:
            print(print_str)

    @staticmethod
    def _draw_progress_bar(total_epoches: int, current_epoch: int, bar_length: int = 30):
        """
        进度条显示
        """
        ratio = current_epoch / total_epoches
        filled = int(ratio * bar_length)
        bar = "█" * filled + "░" * (bar_length - filled)
        print(f"Progress: [{bar}] {current_epoch}/{total_epoches}")

    # ---------------------------------
    # 保存 matplotlib 图片
    def _save_plot(self, epoch=None):
        plt.figure(figsize=self.config.figsize)
        for metric in self.metrics:
            plt.plot(self.show_data[metric], label=metric)
        plt.xlabel("Epoch")
        plt.ylabel("Value")
        plt.title("Training Metrics")
        plt.legend()
        plt.grid(True)

        filename = (
            f"{self.visual_dir}/metrics.png"
            if not self.config.save_each_epoch or epoch is None
            else f"{self.visual_dir}/epoch_{epoch:03d}.png"
        )
        plt.savefig(filename)
        plt.close()

    # ---------------------------------
    # 清屏方法
    # ---------------------------------
    def _clear_terminal(self):
        if not self.config.clear_terminal:
            return

        is_pycharm_console = "PYCHARM_HOSTED" in os.environ

        try:
            if self._have_ipython:
                self._ipython_clear(wait=True)
                return
        except Exception:
            pass

        try:
            if not is_pycharm_console:
                os.system("cls" if os.name == "nt" else "clear")
                return
        except Exception:
            pass

        try:
            sys.stdout.write("\033[2J\033[3J\033[H")
            sys.stdout.flush()
            return
        except Exception:
            pass

        print("\n" * 60)

    # ---------------------------------
    # 保存最终汇总图
    # ---------------------------------
    def finalize(self):
        if not self.config.enable or not self.visual_dir:
            return
        self._save_plot(epoch=None)
        print()  # 确保进度条后换行

    @staticmethod
    def get_show_values(values, value_len):
        """构建终端显示数据"""
        if value_len >= len(values):
            return values
        else:
            data = []
            quotient, remainder = divmod(len(values), value_len)
            for i in range(remainder):
                data.append(values[quotient * i])
            data.extend(values[quotient * remainder:])
            return data


# ========================================
# 测试示例
# ========================================
if __name__ == "__main__":
    vis_config = VisualizationConfig(
        enable=True,
        metrics=["loss", "accuracy"],
        show_terminal_plot=True,
        save_plots=True,
        save_each_epoch=False,
        save_dir="./plots",
        terminal_height=10,
        terminal_width=50,
        terminal_char="●"
    )

    total_epochs = 50
    visualizer = VisualizationManager(vis_config)

    for epoch in range(1, total_epochs + 1):
        loss = 5 / (epoch ** 0.5) + 0.1 * math.sin(epoch)
        accuracy = min(0.9, 0.5 + 0.01 * epoch + 0.05 * math.cos(epoch))
        visualizer.update(total_epochs=total_epochs, loss=loss, accuracy=accuracy)
        time.sleep(0.15)

    visualizer.finalize()
